#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import time

sys.path.append('.')

"""
    手工过滤不符合要求的图片的可视化工具
"""


h_image_count = 10
v_image_count = 10
img_width = 96
img_height = 96


ix, iy = -1, -1
draw_image = None
image_bboxes = []
select_image_bbox_idxes = set()


def get_xy_bbox_idx(x, y):
    global image_bboxes
    bbox_idx = None
    for i, bbox in enumerate(image_bboxes):
        # print(bbox, x, y)
        if (bbox[0] < x < bbox[2]) and (bbox[1] < y < bbox[3]):
            bbox_idx = i
            break
    return bbox_idx


def mouse_event(event, x, y, flags, param):
    global draw_image, image_bboxes, select_image_bbox_idxes

    if draw_image is not None:
        if event == cv2.EVENT_LBUTTONDOWN:
            # print(x, y)
            bbox_idx = get_xy_bbox_idx(x, y)
            if bbox_idx is not None:
                if bbox_idx not in select_image_bbox_idxes:
                    select_image_bbox_idxes.add(bbox_idx)
                else:
                    select_image_bbox_idxes.remove(bbox_idx)


def read_file_list(file_list):
    file_count = 0
    file_select_infos = []
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 1000 == 0:
                print('Read File Count: {}'.format(file_count))

            # gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190914_155747 F0010B00.jpg F0010B01.jpg;0 0
            lines = line.strip().split(';')

            lines0 = lines[0].split()
            lines1 = lines[1].split()
            file_path = lines0[0]
            file_names = lines0[1:]
            selects = list(map(int, lines1))

            file_select_infos.append([file_path, file_names, selects])

        print('Read File Count: {}'.format(file_count))

    return file_select_infos


def save_process(select_file_list_path, file_select_list, cur_img_idx_file_path, cur_img_idx):
    if len(file_select_list) > 0:
        with open(select_file_list_path, 'w') as file:
            for file_select in file_select_list:
                file_dir, file_names, file_selects = file_select
                file_names = ' '.join(file_names)
                file_selects = ' '.join(map(str, file_selects))
                select_info = '{} {};{}\n'.format(file_dir, file_names, file_selects)
                file.write(select_info)
        os.system('chmod a+wr {}'.format(select_file_list_path))

    with open(cur_img_idx_file_path, 'w') as file:
        process_info = '{}\n'.format(cur_img_idx)
        file.write(process_info)
    os.system('chmod a+wr {}'.format(cur_img_idx_file_path))


def filter_select_face(output_dir, file_select_list):
    filter_file_list_path = os.path.join(output_dir, 'FilterSelectImage.txt')
    filter_file_list_file = open(filter_file_list_path, 'w')
    face_count = 0
    face_image_count = 0
    file_select_list.sort(key=lambda x: x[0])
    for dir_name, files, selects in file_select_list:
        file_names = []
        for file_name, is_select in zip(files, selects):
            if 0 != is_select:
                continue
            file_names.append(file_name)
        if len(file_names) > 0:
            file_names.sort()
            for file_name in file_names:
                out_info = '{}/{}\n'.format(dir_name, file_name)
                filter_file_list_file.write(out_info)
                face_image_count += 1
            face_count += 1
    filter_file_list_file.close()
    os.system('chmod a+wr {}'.format(filter_file_list_path))
    print('Face Count: {}, Face Image Count: {}'.format(face_count, face_image_count))


def read_bboxes(bbox_count_list_path, bbox_list_path, score_thr):
    dir_file_selects = []
    scores_list = []
    with open(bbox_count_list_path, 'r') as bbox_count_file:
        with open(bbox_list_path, 'r') as bbox_file:
            for line in bbox_count_file.readlines():
                lines = line.strip().split(',')
                file_dir = os.path.splitext(lines[0])[0]

                bbox_count = int(lines[1])
                file_names = []
                scores = []
                selects = []
                for _ in range(bbox_count):
                    bbox_infos = bbox_file.readline()
                    bbox_infos = bbox_infos.strip().split()
                    file_name = bbox_infos[0]
                    bbox = bbox_infos[1].split(',')
                    score = float(bbox[-1])
                    file_names.append(file_name)
                    scores.append(score)
                    sel = 1 if score < score_thr else 0
                    selects.append(sel)
                dir_file_selects.append([file_dir, file_names, selects])
                scores_list.append(scores)
    return dir_file_selects, scores_list


def main():
    global draw_image, image_bboxes, select_image_bbox_idxes
    parser = argparse.ArgumentParser(description="Frame Face Select")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data1/data/FaceVideo/CropImageMargin", help="path to image file", type=str
    )
    parser.add_argument(
        "--bbox_count_list", default="/rootfs/media/kasim/Data1/data/FaceVideo/CropImageMargin/bbox_count_list.txt", help="path to bbox count list", type=str
    )
    parser.add_argument(
        "--bbox_list", default="/rootfs/media/kasim/Data1/data/FaceVideo/CropImageMargin/bbox_list.txt", help="path to bbox list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/FaceVideo/CropImageMargin", help="path to mot reid image info", type=str
    )
    parser.add_argument("--use_temp", action='store_true', help="use temp save, if crash can use it")
    parser.add_argument('--save_time', type=int, default=60, help='save time, second ')
    parser.add_argument("--filter", action='store_true', help="filter select image")
    parser.add_argument('--score_thr', type=float, default=0.2, help='bbox score thr')
    args = parser.parse_args()

    save_time = args.save_time
    image_dir = args.image_dir
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    display_image_count = h_image_count * v_image_count

    select_file_list_path = os.path.join(output_dir, 'select_images.txt')
    file_select_list, scores_list = read_bboxes(args.bbox_count_list, args.bbox_list, args.score_thr)
    if os.path.exists(select_file_list_path):
        file_select_list = read_file_list(select_file_list_path)

    if args.filter:
        filter_select_face(output_dir, file_select_list)
        return

    cv2.namedWindow('image')
    cv2.setMouseCallback('image', mouse_event)

    cur_img_idx = 0
    cur_img_idx_file_path = os.path.join(output_dir, 'cur_img_idx.txt')
    if os.path.exists(cur_img_idx_file_path):
        with open(cur_img_idx_file_path, 'r') as file:
            for line in file.readlines():
                cur_img_idx = int(line.strip())
                break

    if args.use_temp:
        tmp_output_dir = os.path.join('/dev/shm', 'SelectImage')
        if not os.path.exists(tmp_output_dir):
            os.makedirs(tmp_output_dir)
            os.system('chmod a+wr {}'.format(tmp_output_dir))

        tmp_select_file_list_path = os.path.join(tmp_output_dir, 'select_images.txt')
        tmp_cur_id_idx_file_path = os.path.join(tmp_output_dir, 'cur_img_idx.txt')

    select_change = False
    start_save_time = time.time()

    last_idx = len(file_select_list) - 1

    try:
        while True:
            idx = cur_img_idx
            if idx >= last_idx:
                break
            file_dir, file_names, file_selects = file_select_list[idx]
            scores = scores_list[idx]

            start_idx = 0
            files = []

            _last_idx = len(file_names) - 1
            for _idx, file_name in enumerate(file_names):
                files.append(file_name)
                file_count = len(files)
                if file_count == 1:
                    start_idx = _idx
                if file_count < display_image_count and (_idx < _last_idx):
                    continue
                imgs = []
                for file_name in files:
                    file_path = os.path.join(image_dir, file_dir, file_name)
                    img = cv2.imread(file_path)
                    # img = cv2.resize(img, (img_width, img_height))
                    w, h = img.shape[1], img.shape[0]
                    wscale = img_width / w
                    hscale = img_height / h
                    scale = min(wscale, hscale)
                    w = int(round(w * scale))
                    h = int(round(h * scale))
                    image = cv2.resize(img, (w, h))
                    img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
                    x = (img_width-w)//2
                    y = (img_height-h)//2
                    img[y:y+h, x:x+w, :] = image
                    imgs.append(img)

                image_count = len(imgs)
                images = []
                x_off = 0
                y_off = 0
                image_bboxes.clear()
                select_image_bbox_idxes.clear()
                for select_idx in range(0, image_count):
                    if 0 != file_selects[start_idx + select_idx]:
                        select_image_bbox_idxes.add(select_idx)
                for j, img in enumerate(imgs):
                    if j % h_image_count == 0:
                        x_off = 0
                        if j != 0:
                            y_off += img_height
                        h_images = [img]
                    else:
                        x_off += img_width
                        h_images.append(img)
                    image_bboxes.append([x_off, y_off, (x_off + img_width), (y_off + img_height)])
                    if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                        images.append(np.concatenate(h_images, axis=1))

                if len(images) > 1:
                    w0 = images[0].shape[1]
                    h0 = images[0].shape[0]
                    w1 = images[-1].shape[1]
                    h1 = images[-1].shape[0]
                    if (w0 != w1) or (h0 != h1):
                        # image = np.empty_like(images[0])
                        image = np.zeros_like(images[0])
                        image[0:h1, 0:w1, :] = images[-1]
                        images[-1] = image
                draw_image = np.vstack(images)

                change_trk_id = False
                is_exit = False
                while True:
                    image = draw_image.copy()
                    for select_idx in select_image_bbox_idxes:
                        bbox = image_bboxes[select_idx]
                        cv2.circle(image, (bbox[0] + 14, bbox[1] + 14), 8, (0, 0, 255), -1)
                        # cv2.circle(image, (bbox[0] + 14, bbox[3] - 14), 8, (0, 0, 255), -1)

                    for bbox_idx in range(image_count):
                        select_idx = start_idx + bbox_idx
                        bbox = image_bboxes[bbox_idx]
                        score = scores[select_idx]
                        cv2.putText(image, '{:.03f}'.format(score), (bbox[0] + 8, bbox[3] - 8),
                                    cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=2)

                    cv2.putText(image, '%d/%d %d/%d' % (cur_img_idx + 1, last_idx + 1, _idx + 1, _last_idx + 1),
                                (8, 24), cv2.FONT_HERSHEY_PLAIN, 2.0, (0, 255, 0), thickness=2)
                    cv2.imshow('image', image)
                    key = cv2.waitKey(50)
                    if (key == 27) or (key == ord('q')) or (key == ord('Q')):
                        is_exit = True
                        break
                    elif key == ord(' '):  # 同一ID的下一页
                        select_change = True
                        break
                    elif key == ord('a') or key == ord('A'):  # 上一个ID
                        select_change = True
                        if cur_img_idx > 0:
                            cur_img_idx -= 1
                            change_trk_id = True
                        break
                    elif key == ord('d') or key == ord('D'):  # 下一个ID
                        select_change = True
                        change_trk_id = True
                        cur_img_idx += 1
                        break
                    elif key == ord('s') or key == ord('S'):  # Save
                        if args.use_temp:
                            save_process(tmp_select_file_list_path, file_select_list, tmp_cur_id_idx_file_path, cur_img_idx)
                        start_save_time = time.time()
                        select_change = False
                        break
                    if args.use_temp:
                        # 定时保存
                        if time.time() - start_save_time >= save_time:
                            if select_change:
                                save_process(tmp_select_file_list_path, file_select_list, tmp_cur_id_idx_file_path, cur_img_idx)
                                select_change = False
                            start_save_time = time.time()
                if len(select_image_bbox_idxes) > 0:
                    for bbox_idx in range(0, image_count):
                        select_idx = start_idx + bbox_idx
                        if bbox_idx in select_image_bbox_idxes:
                            file_selects[select_idx] = 1
                        else:
                            file_selects[select_idx] = 0
                files.clear()
                if is_exit:
                    save_process(select_file_list_path, file_select_list, cur_img_idx_file_path, cur_img_idx)
                    return
                if change_trk_id:
                    break

    except Exception as e:
        print(e)
    save_process(select_file_list_path, file_select_list, cur_img_idx_file_path, cur_img_idx)
    print('finish!')


if __name__ == '__main__':
    main()
